package dubrouski.csc;

import dubrouski.io.FileWriter;
import dubrouski.segmentation.CentroidsCalculator;
import dubrouski.segmentation.CentroidsCalculatorImpl;
import dubrouski.triangulation.Triangle;
import dubrouski.triangulation.Triangulate;
import dubrouski.triangulation.Vector3d;
import dubrouski.util.Utils;
import ij.ImagePlus;
import ij.ImageStack;
import ij.gui.GenericDialog;
import ij.plugin.filter.PlugInFilter;
import ij.process.ImageProcessor;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import org.la4j.LinearAlgebra;
import org.la4j.linear.LinearSystemSolver;
import org.la4j.matrix.dense.Basic1DMatrix;
import org.la4j.vector.dense.BasicVector;

/**
 *
 * @author Stanislau Dubrouski
 */
public class CSC implements PlugInFilter {

    private ImageStack stack;
    static final int X = 0;
    static final int Y = 1;
    static final int Z = 2;
    /*
     * Params for wrong values elimination. 
     * Values below are taken from 
     * 'An efficient algorithm for measurement 
     * and correction of chromatic aberrations 
     * in fluorescence microscopy.'
     */ 
    static float FIRST_PASS = 4;
    static float SECOND_PASS = 3;

    /**
     *
     * @param arg
     * @param imp
     * @return
     */
    public int setup(String arg, ImagePlus imp) {
        if (imp != null){
            this.stack = imp.getStack();   
        }
        return DOES_RGB + STACK_REQUIRED;
    }

    /**
     *
     * @param ip
     */
    public void run(ImageProcessor ip) {

        UsersParameters usersParams = showDialogWithParams();
        if (usersParams == null) {
            return;
        }
        /*
         * Allow users set their own param values.
         */
        FIRST_PASS = usersParams.getFpass();
        SECOND_PASS = usersParams.getSpass();

        List<Colors> colors = getProcessedColors(usersParams);

        //start time counter
        long veryBeginning = System.currentTimeMillis();
        //[x][y][z]
        int[][][] image = CSC.convertTo3DArray(stack);

        Thresholder t = new BasicThresholderImpl();
        image = t.thresholdImage(image, usersParams.getThresholdLevel());
        
        CentroidsCalculator centroidsCalc = new CentroidsCalculatorImpl();
        List<List<Vector3d>> spots = centroidsCalc.computeCentroids(colors, image);

        spots = associateSpots(spots, usersParams.getMaxCentroidDistance(), colors);

        //compute shift for every colors pair
        List<List<Shift>> ShiftsForColorPairs = createShiftVectorsForColorPairs(spots, colors);
        //PrintVectors(ShiftsForColorPairs, "_before_filtering");

        ShiftsForColorPairs = filterShifts(ShiftsForColorPairs, FIRST_PASS);        
        //PrintVectors(ShiftsForColorPairs, "_first_pass");

        ShiftsForColorPairs = filterShifts(ShiftsForColorPairs, SECOND_PASS);
        PrintVectors(ShiftsForColorPairs, "_second_pass");

        List<List<PlaneRepresentor>> planes = calculatePlanes(ShiftsForColorPairs);
        PrintPlaneEquations(planes);
        
        //create mock shifts to visualize shifts maps
        List<List<Shift>> calcShifts = modelShiftsBasedOnPlanesEquations(ShiftsForColorPairs, planes);

        //print mocks to form shifts visualization
        PrintVectors(calcShifts, "_approximate_shifts");
        
        long endTime = System.currentTimeMillis();
        long elapsedTime = endTime - veryBeginning;

        
        
        
        
        
        
        
        
        /*
         * TODO
         * testing correcting
         * repeat computing, then correct computed positions using planes equations
        */
        //[x][y][z]
        int[][][] image2 = CSC.convertTo3DArray(stack);

        Thresholder t2 = new BasicThresholderImpl();
        image2 = t2.thresholdImage(image2, usersParams.getThresholdLevel());
        
        CentroidsCalculator centroidsCalc2 = new CentroidsCalculatorImpl();
        List<List<Vector3d>> spots2 = centroidsCalc2.computeCentroids(colors, image2);

        spots2 = associateSpots(spots2, usersParams.getMaxCentroidDistance(), colors);

        //correct spots positions
        spots2 = applyCorrectionOnSpotsPositions(spots2, planes);
        
        //compute shift for every colors pair
        List<List<Shift>> ShiftsForColorPairs2 = createShiftVectorsForColorPairs(spots2, colors);
        
        
        PrintVectors(ShiftsForColorPairs2, "_corrected_spots");

        
        
        
        
        
        
        
        
        
        
        showResultsDialog(elapsedTime);
    }
    
    
    private List<List<Vector3d>> applyCorrectionOnSpotsPositions(List<List<Vector3d>> spots, List<List<PlaneRepresentor>> planes){
        List<List<Vector3d>> result = new ArrayList<List<Vector3d>>(3);
        
        
        for (int i = 0; i < spots.size(); i++){
            result.add(applyCorrection(spots.get(i), planes.get(i)));            
        }
        
        
        
        return result;
    }
    
    private List<Vector3d> applyCorrection(List<Vector3d> spots, List<PlaneRepresentor> planes){
        List<Vector3d> result = new ArrayList<Vector3d>();
        
        for (Vector3d v : spots){
            Vector3d correctedV = applyCorrectionOnX(v, planes.get(0));
            correctedV = applyCorrectionOnY(correctedV, planes.get(1));
            result.add(correctedV);
        }        
        return result;
    }
    
    private Vector3d applyCorrectionOnX(Vector3d spot, PlaneRepresentor plane){
        if (spot == null){
            return null;
        }
        Vector3d result = new Vector3d(spot.getX() + plane.calculateZ(spot.getX(), spot.getY()), spot.getY(), spot.getZ());        
        return result;
    }
    
    private Vector3d applyCorrectionOnY(Vector3d spot, PlaneRepresentor plane){
        if (spot == null){
            return null;
        }
        
        Vector3d result = new Vector3d(spot.getX(), spot.getY() + plane.calculateZ(spot.getX(), spot.getY()), spot.getZ());        
        return result;
    }

    /*
     * Calculates planes equations based on given shift values.
     */
    private List<List<PlaneRepresentor>> calculatePlanes(List<List<Shift>> ShiftsForColorPairs) {
        /*
         * Left-side part of equation - spots coordinates x and y (1st and 2nd columns, 2rd column filled with 0s).
         * There is a separate matrix for every color pair.
         */
        List<Basic1DMatrix> leftSideMatrixes = new ArrayList<Basic1DMatrix>();

        for (int i = 0; i < ShiftsForColorPairs.size(); i++) {
            leftSideMatrixes.add(new Basic1DMatrix(ShiftsForColorPairs.get(i).size(), 3));
            for (int j = 0; j < ShiftsForColorPairs.get(i).size(); j++) {

                leftSideMatrixes.get(i).setRow(j,
                        new BasicVector(new double[]{
                    ShiftsForColorPairs.get(i).get(j).getBeginVector().getX(),
                    ShiftsForColorPairs.get(i).get(j).getBeginVector().getY(),
                    1
                }));
            }
        }


        //TODO REMOVE
        System.err.println("Left side matrix 1");
        System.err.println(leftSideMatrixes.get(0).toString());
        System.err.println(leftSideMatrixes.get(1).toString());
        System.err.println(leftSideMatrixes.get(2).toString());
        System.err.println("End of Left side matrix 1");

        /*
         * Right-side part of equation - shift value for every component (x, y, z).
         * There is a separate vector for every color pair.
         */
        ArrayList<List<BasicVector>> rightSideVectors = new ArrayList<List<BasicVector>>();

        for (int i = 0; i < ShiftsForColorPairs.size(); i++) {
            rightSideVectors.add(new ArrayList<BasicVector>(ShiftsForColorPairs.get(i).size()));
            /*
             * Add BasicVector for every shift component (x, y, z).
             */
            rightSideVectors.get(i).add(new BasicVector(
                    ShiftsForColorPairs.get(i).size()));
            rightSideVectors.get(i).add(new BasicVector(
                    ShiftsForColorPairs.get(i).size()));
            rightSideVectors.get(i).add(new BasicVector(
                    ShiftsForColorPairs.get(i).size()));

            for (int j = 0; j < ShiftsForColorPairs.get(i).size(); j++) {
                //set x-component of the shift
                rightSideVectors.get(i).get(0).set(j,
                        Shift.computeVectorValue(ShiftsForColorPairs.get(i).get(j)).getX());
                //set y-component of the shift
                rightSideVectors.get(i).get(1).set(j,
                        Shift.computeVectorValue(ShiftsForColorPairs.get(i).get(j)).getY());
                //set z-component of the shift
                rightSideVectors.get(i).get(2).set(j,
                        Shift.computeVectorValue(ShiftsForColorPairs.get(i).get(j)).getZ());
            }
        }

        //TODO REMOVE
        System.err.println("Right side matrix 1");
        System.err.println(rightSideVectors.get(0).get(0));
        System.err.println(rightSideVectors.get(0).get(1));
        System.err.println(rightSideVectors.get(1).get(0));
        System.err.println(rightSideVectors.get(1).get(1));
        System.err.println(rightSideVectors.get(2).get(0));
        System.err.println(rightSideVectors.get(2).get(1));
        System.err.println("End of Right side matrix 1");

        /*
         * Calculated planes for 
         * Each shift component in each colorpair is represented by separate plane.
         */
        List<List<PlaneRepresentor>> planes = new ArrayList<List<PlaneRepresentor>>();

        for (int colorPair = 0; colorPair < ShiftsForColorPairs.size(); colorPair++) {
            planes.add(new ArrayList<PlaneRepresentor>());
            LinearSystemSolver solver = leftSideMatrixes.get(colorPair).withSolver(LinearAlgebra.LEAST_SQUARES);
            //TODO now calculates also z-component. Delete, because z is calced with different equation.
            for (int shiftComponent = 0; shiftComponent < 3; shiftComponent++) {
                org.la4j.vector.Vector res = solver.solve(rightSideVectors.get(colorPair).get(shiftComponent), LinearAlgebra.DENSE_FACTORY);
                planes.get(colorPair).add(new PlaneRepresentor(res.get(0), res.get(1), res.get(2)));
            }
        }

        return planes;

    }

    private List<Vector3d> getBeginVectorsListFromShifts(List<Shift> input) {
        List<Vector3d> result = new ArrayList<Vector3d>();
        for (Shift pair : input) {
            result.add(new Vector3d(pair.getBeginVector().getX(), pair.getBeginVector().getY(), pair.getBeginVector().getZ()));
        }
        return result;
    }

    private List<List<Shift>> filterShifts(List<List<Shift>> input, float coeff) {
        for (List<Shift> shifts : input) {
            List<Vector3d> spots = getBeginVectorsListFromShifts(shifts);
            List<Triangle> triangles = Triangulate.triangulate((ArrayList<Vector3d>) Utils.removeNullValues(spots));

            List<Shift> wrongVals = new ArrayList<Shift>();
            wrongVals = getWrongShiftValues(shifts, triangles);

            shifts.removeAll(wrongVals);
        }
        return input;
    }

    private List<Shift> getWrongShiftValues(
            List<Shift> pairs, List<Triangle> triangles) {

        List<Vector3d> averageNeighbourVectors = new ArrayList<Vector3d>();
        for (Shift pair : pairs) {
            averageNeighbourVectors.add(getAverageNeighbourVector(pair, pairs, triangles));
        }

        List<Float> diffs = new ArrayList<Float>();
        for (int i = 0; i < pairs.size(); i++) {
            Vector3d vectorValue = Shift.computeVectorValue(pairs.get(i));
            diffs.add(new Float(Vector3d.computeNormVector(
                    Vector3d.substractVectors(vectorValue, averageNeighbourVectors.get(i)))));
        }

        float meanDiff = Utils.computeMeanValue(diffs);
        float deviation = Utils.computeStandartDeviation(diffs);

        List<Shift> wrongValues = new ArrayList<Shift>();
        float diff;
        float measure = FIRST_PASS * deviation;
        for (int i = 0; i < pairs.size(); i++) {
            diff = Math.abs(diffs.get(i) - meanDiff);
            if (diff > measure) {
                wrongValues.add(new Shift(pairs.get(i)));
            }
        }

        //TODO remove test draw found wrong values stage 1
        FileWriter.writeToFile("_WrongValues.html", Visualizer.drawLCAVectors(wrongValues, "", 100));

        return wrongValues;
    }

    private Vector3d getAverageNeighbourVector(Shift pair, List<Shift> pairs, List<Triangle> triangles) {
        List<Shift> neighbours = getNeighbourPairsOfVectors(pair, pairs, triangles);
        return computeMeanVector(pair, neighbours);
    }

    //TODO public for tests
    /**
     *
     * @param pair
     * @param input
     * @return
     */
    public static Vector3d computeMeanVector(Shift pair, Collection<Shift> input) {
        Vector3d weightedVectorsSum = new Vector3d(0, 0, 0);
        float weightsSum = 0;

        for (Shift neighbour : input) {
            Vector3d v = Shift.computeVectorValue(neighbour);
            float distance = Vector3d.computeVectorsDistance(pair.getBeginVector(), neighbour.getBeginVector());
            float weight = 1 / distance;

            weightedVectorsSum.add(Vector3d.multiplyVector(v, weight));
            weightsSum += weight;
        }
        return Vector3d.divideVector(weightedVectorsSum, weightsSum);
    }

    private List<Shift> getNeighbourPairsOfVectors(Shift pair, List<Shift> pairsOfVectors, List<Triangle> triangles) {
        List<Shift> result = new ArrayList<Shift>();

        Collection<Vector3d> neighbourSpots = getNeighbourSpots(pair.getBeginVector(), triangles);

        for (Shift pv : pairsOfVectors) {
            if (neighbourSpots.contains(pv.getBeginVector())) {
                result.add(pv);
            }
        }

        return result;
    }

    private Collection<Vector3d> getNeighbourSpots(Vector3d pv, List<Triangle> triangles) {
        Set<Vector3d> result = new HashSet<Vector3d>();
        List<Triangle> trianglesWithSpot = getTrianglesWithSpot(pv, triangles);

        for (Triangle tr : trianglesWithSpot) {
            result.add(new Vector3d(tr.getP1().getX(), tr.getP1().getY(), tr.getP1().getZ()));
            result.add(new Vector3d(tr.getP2().getX(), tr.getP2().getY(), tr.getP2().getZ()));
            result.add(new Vector3d(tr.getP3().getX(), tr.getP3().getY(), tr.getP3().getZ()));
        }

        result.remove(pv);
        return result;
    }
    //TODO public for tests

    /**
     *
     * @param vector
     * @param triangles
     * @return
     */
    public static List<Triangle> getTrianglesWithSpot(Vector3d vector, List<Triangle> triangles) {
        List<Triangle> result = new ArrayList<Triangle>();

        for (Triangle tr : triangles) {
            if (triangleContainsVector(vector, tr)) {
                result.add(tr);
                //result.add(new Triangle(tr.getP1(), tr.getP2(), tr.getP3()));
            }
        }
        return result;
    }

    //TODO public for tests
    /**
     *
     * @param pv
     * @param tr
     * @return
     */
    public static boolean triangleContainsVector(Vector3d pv, Triangle tr) {
        return tr.getP1().equals(pv)
                || tr.getP2().equals(pv)
                || tr.getP3().equals(pv);
    }

    private List<Colors> getProcessedColors(UsersParameters usersParams) {
        List<Colors> result = new ArrayList<Colors>();
        if (usersParams.getProcessRed()) {
            result.add(Colors.RED);
        }
        if (usersParams.getProcessGreen()) {
            result.add(Colors.GREEN);
        }
        if (usersParams.getProcessBlue()) {
            result.add(Colors.BLUE);
        }

        return result;
    }

    private int[] getImage3dSize(ImageStack stack) {
        int[] imageSize = new int[3];
        imageSize[X] = stack.getWidth();
        imageSize[Y] = stack.getHeight();
        imageSize[Z] = stack.getSize();
        return imageSize;
    }
    //TODO public for tests

    /**
     *
     * @param stack
     * @return
     */
    public static int[][][] convertTo3DArray(ImageStack stack) {
        int xSize = stack.getWidth();
        int ySize = stack.getHeight();
        int zSize = stack.getSize();
        int[][][] resultImage = new int[stack.getWidth()][stack.getHeight()][stack.getSize()];

        for (int z = 0; z < zSize; z++) {
            int[] pixels = (int[]) stack.getPixels(z + 1);
            for (int x = 0; x < xSize; x++) {
                for (int y = 0; y < ySize; y++) {
                    resultImage[x][y][z] = pixels[y * xSize + x];
                    //TEST - paint threshold results
                    //pixels[y * xSize + x] = image[x][y][z];
                }
            }
        }
        return resultImage;
    }

    private UsersParameters showDialogWithParams() {
        UsersParameters params = new UsersParameters();
        GenericDialog gd = new GenericDialog("Set params");
        gd.addMessage("Select colors:");
        gd.addCheckbox("Red", true);
        gd.addCheckbox("Green", true);
        gd.addCheckbox("Blue", true);
        gd.addNumericField("Set threshold level: ", 120, 0);
        gd.addNumericField("Maximum distance between centroids (px): ", 10, 0);
        gd.addNumericField("Coeff. for first cycle: ", 2.5, 0);
        gd.addNumericField("Coeff. for second cycle: ", 2.5, 0);
        //gd.addStringField("Output to:", "C:\\Users\\Stadu\\Documents\\data.txt", 50);
        gd.showDialog();
        if (gd.wasCanceled()) {
            return null;
        }
        params.setProcessRed(gd.getNextBoolean());
        params.setProcessGreen(gd.getNextBoolean());
        params.setProcessBlue(gd.getNextBoolean());
        int thresholdLevel = (int) gd.getNextNumber();
        params.setThresholdLevel(thresholdLevel);
        int maxDistance = (int) gd.getNextNumber();
        params.setMaxCentroidDistance(maxDistance);
        float fPass = (float) gd.getNextNumber();
        params.setFpass(fPass);
        float sPass = (float) gd.getNextNumber();
        params.setSpass(sPass);

        //String outputTo = gd.getNextString();

        if (!params.getProcessRed() && !params.getProcessGreen() && !params.getProcessBlue()) {
            GenericDialog error = new GenericDialog("Error");
            error.addMessage("No colors chosen. Nothing to do.");
            error.showDialog();
            return null;
        }
        return params;
    }

    private void showText(String text) {
        GenericDialog results = new GenericDialog("Here are some results");
        results.addMessage(text);
        results.showDialog();
    }

    private void showResultsDialog(long elapsedTime) {
        GenericDialog results = new GenericDialog("Here are some results");
        results.addMessage("Execution took " + elapsedTime + " ms.");
        results.showDialog();
    }

    private List<List<Vector3d>> associateSpots(List<List<Vector3d>> spots, int maxCentroidDistance, List<Colors> colors) {

        if (colors.size() == 1) {
            return spots;
        }

        List<List<Vector3d>> result = new ArrayList<List<Vector3d>>();

        for (int color = 0; color < colors.size(); color++) {
            result.add(new ArrayList<Vector3d>());
        }

        while (true) {
            int spotTakenFrom = -1;

            for (int color = 0; color < colors.size(); color++) {
                if (!spots.get(color).isEmpty()) {
                    result.get(color).add(spots.get(color).get(0));
                    spots.get(color).remove(0);
                    spotTakenFrom = color;
                    break;
                }
            }

            if (spotTakenFrom < 0) {
                break;
            }

            //insert nulls for already emptied colors
            for (int color = 0; color < spotTakenFrom; color++) {
                result.get(color).add(null);
            }

            //find closest spots for every not-empty color
            for (int color = spotTakenFrom + 1; color < colors.size(); color++) {
                Vector3d closestSpot = findClosestSpot(result.get(spotTakenFrom).get(result.get(spotTakenFrom).size() - 1), spots.get(color), maxCentroidDistance);
                result.get(color).add(closestSpot);
                spots.get(color).remove(closestSpot);
            }
        }

        return result;
    }

    private Vector3d findClosestSpot(Vector3d spot, List<Vector3d> spots, int maxCentroidDistance) {
        Vector3d result = null;

        if (spots.isEmpty()) {
            return null;
        }

        float maxDistance = 1000;

        for (int i = 0; i < spots.size(); i++) {
            if ((Vector3d.computeVectorsDistance(spot, spots.get(i)) < maxCentroidDistance) && (Vector3d.computeVectorsDistance(spot, spots.get(i)) < maxDistance)) {
                maxDistance = Vector3d.computeVectorsDistance(spot, spots.get(i));
                result = spots.get(i);
            }
        }
        return result;
    }

    private List<List<Shift>> createShiftVectorsForColorPairs(List<List<Vector3d>> pVectorLists, List<Colors> colors) {
        List<List<Shift>> shiftVectors = new ArrayList<List<Shift>>();

        //1. color to 2. color shift
        shiftVectors.add(new ArrayList<Shift>());
        for (int point = 0; point < pVectorLists.get(0).size(); point++) {
            if (pVectorLists.get(0).get(point) == null || pVectorLists.get(1).get(point) == null) {
                continue;
            }
            shiftVectors.get(0).add(new Shift(pVectorLists.get(0).get(point), pVectorLists.get(1).get(point)));
        }

        if (colors.size() == 3) {
            //2. color to 3. color shift
            shiftVectors.add(new ArrayList<Shift>());
            for (int point = 0; point < pVectorLists.get(1).size(); point++) {
                if (pVectorLists.get(1).get(point) == null || pVectorLists.get(2).get(point) == null) {
                    continue;
                }
                shiftVectors.get(1).add(new Shift(pVectorLists.get(1).get(point), pVectorLists.get(2).get(point)));
            }

            //3. color to 1. color shift
            shiftVectors.add(new ArrayList<Shift>());
            for (int point = 0; point < pVectorLists.get(2).size(); point++) {
                if (pVectorLists.get(2).get(point) == null || pVectorLists.get(0).get(point) == null) {
                    continue;
                }
                shiftVectors.get(2).add(new Shift(pVectorLists.get(2).get(point), pVectorLists.get(0).get(point)));
            }
        }

        return shiftVectors;
    }

    private void PrintVectors(List<List<Shift>> ShiftsForColorPairs, String name) {
        for (int i = 0; i < ShiftsForColorPairs.size(); i++) {
            FileWriter.writeToFile(i + name + ".html", Visualizer.drawLCAVectors(ShiftsForColorPairs.get(i), "", 100));
        }
    }

    private void PrintPlaneEquations(List<List<PlaneRepresentor>> planes) {
        //TODO Remove printing
        System.err.println("========================");

        System.err.println("R->G, x-component:");
        System.err.println(planes.get(0).get(0));

        System.err.println("G->B, x-component:");
        System.err.println(planes.get(1).get(0));

        System.err.println("B->R, x-component:");
        System.err.println(planes.get(2).get(0));

        System.err.println("R->G, y-component:");
        System.err.println(planes.get(0).get(1));

        System.err.println("G->B, y-component:");
        System.err.println(planes.get(1).get(1));

        System.err.println("B->R, y-component:");
        System.err.println(planes.get(2).get(1));

        System.err.println("========================");
    }

    private List<List<Shift>> modelShiftsBasedOnPlanesEquations(List<List<Shift>> ShiftsForColorPairs, List<List<PlaneRepresentor>> planes) {
        /*
         * build mock shifts based on calculated planes to visualize approximated shifts map.
         */
        List<List<Shift>> calcShifts = new ArrayList<List<Shift>>();
        for (int colorPair = 0; colorPair < ShiftsForColorPairs.size(); colorPair++) {
            calcShifts.add(new ArrayList<Shift>());
            for (int i = 0; i < 1000; i += 30) {
                for (int j = 0; j < 1000; j += 30) {
                    calcShifts.get(colorPair).add(
                            new Shift(
                            new Vector3d(i, j, 0),
                            new Vector3d(i + planes.get(colorPair).get(0).calculateZ(i, j), j + planes.get(colorPair).get(1).calculateZ(i, j), 0)));
                }
            }
        }
        return calcShifts;
    }
}